//
//  ukf.c
//  KFilter
//
//  Created by Kaiwen Gu on 6/23/13.
//  Copyright (c) 2013 com.mika. All rights reserved.
//

#include <stdio.h>
#include <stdlib.h>
#include "ukf.h"
#include "erza_matrix_lib.h"
#include "config_default.h"

int setCDKFDiagonalCovariance(arm_matrix_instance_f32 * instance, uint16_t rcDim, float32_t * cov, uint16_t covDim, uint8_t doSqrt) {
	int i;

	if (rcDim != covDim) {
		return 0;
	}

	matrixInit(instance, rcDim, rcDim);

	if (!instance) {
		return 0;
	}

	for (i = 0; i < rcDim; i++) {
			instance->pData[i + i * rcDim] = sqrtf(cov[i]);
	}
	return 1;
}

/*Init the UKF*/
int ukf_Init(ukf_predictionModel *model) {
	int sigmaPointsDim = 1 + 2 * (model->stateDim + model->observationCovDim);
	int augmentedDim = (model->stateDim + model->observationCovDim);

	model->Sx = model->processCov;
	model->processCovDim = model->stateDim;
	/*
	 W1 = [(hh - Xdim - Vdim)/hh   1/(2*hh);                  % sigma-point weights set 1
	 1/(2*h)                 sqrt(hh-1)/(2*hh)];
	 */
	model->scaleFactorSqrt = sqrtf(model->scaleFactor);
	model->weightSum0 = (model->scaleFactor - model->processCovDim - model->observationCovDim) / model->scaleFactor;
	model->weightSum1 = 1 / (2 * model->scaleFactor);
	model->weightCov0 = 1 / (2 * model->scaleFactorSqrt);
	model->weightCov1 = sqrtf(model->scaleFactor - 1) / (4 * model->scaleFactor);

	/*State matrix*/
	matrixInit(&(model->state), model->stateDim, ONE_CR);

	/*Sigma point augmented with covariance*/
	matrixInit(&(model->hSzMT), sigmaPointsDim, augmentedDim);

	/*propagated state equations*/
	matrixInit(&(model->X), sigmaPointsDim, (model->stateDim));

	/*newly calc'ed cov variance*/
	matrixInit(&(model->Sxn), (model->stateDim), 2 * (model->stateDim + model->observationCovDim));
	matrixInit(&(model->SxnT), (model->stateDim), (model->stateDim));

	/*Observation*/
	matrixInit(&(model->Y), MAX_OBS_LENGTH, 1 + 2 * ((model->stateDim) + MAX_OBS_LENGTH));
	matrixInit(&(model->yh), MAX_OBS_LENGTH, ONE_CR);

	matrixInit(&(model->Syn), MAX_OBS_LENGTH, 2 * (model->stateDim + MAX_OBS_LENGTH));
	matrixInit(&(model->SynT), MAX_OBS_LENGTH, MAX_OBS_LENGTH);
	matrixInit(&(model->Sy), MAX_OBS_LENGTH, MAX_OBS_LENGTH);

	matrixInit(&(model->Syx), model->stateDim, MAX_OBS_LENGTH);
	matrixInit(&(model->SyxT), MAX_OBS_LENGTH, model->stateDim);

	matrixInit(&(model->Syw), MAX_OBS_LENGTH, MAX_OBS_LENGTH);
	matrixInit(&(model->Pxy), MAX_OBS_LENGTH, model->stateDim);
	matrixInit(&(model->KG), MAX_OBS_LENGTH, model->stateDim);

	matrixInit(&(model->innov), MAX_OBS_LENGTH, ONE_CR);
	matrixInit(&(model->innovUpdate), MAX_OBS_LENGTH, ONE_CR);

	matrixInit(&(model->s1), model->stateDim, model->stateDim + MAX_OBS_LENGTH);	// scratch
	matrixInit(&(model->s2), model->stateDim, model->stateDim);	// scratch
	matrixInit(&(model->s3), model->stateDim, MAX_OBS_LENGTH);	// scratch
	matrixInit(&(model->s4), model->stateDim, model->stateDim + MAX_OBS_LENGTH);	// scratch
	matrixInit(&(model->s5), 2 * (model->stateDim + MAX_OBS_LENGTH), model->stateDim);	// scratch
	matrixInit(&(model->s5T), model->stateDim, 2 * (model->stateDim + MAX_OBS_LENGTH));	// scratch

	matrixInit(&(model->D), model->stateDim + MAX_OBS_LENGTH, MAX_OBS_LENGTH);	// D matrix

	return 1;
}

/*Transposed sigma points, easier to propagate through models*/
int getTransposedSigmaPoints(ukf_predictionModel * model, arm_matrix_instance_f32 processNoiseCov) {
	int row, col, arrIndex, t1, t2, t3;
	float32_t stateElement, covarianceElement;

	/*dimensions*/
	const float32_t scaleFactor = model->scaleFactorSqrt;

	const int stateDim = model->stateDim;
	const int processCovDim = model->processCovDim;
	const int processCovDimAdd = processCovDim + 1;
	const int processNoiseCovDim = processNoiseCov.numCols;
	const int augmentedDim = processCovDim + processNoiseCovDim;
	const int augmentedDimAdd = augmentedDim + 1;
	const int numSigmaPoints = 1 + (augmentedDim << 1);

	/*data arrays*/
	const float32_t * state = model->state.pData;
	const float32_t * predictedCov = model->Sx.pData;
	const float32_t * processCov = processNoiseCov.pData;

	float32_t * sigmaPointsArray = model->hSzMT.pData;
	float32_t * sigmaPointsArrayAug = sigmaPointsArray + augmentedDim * augmentedDim;

	model->hSzMT.numCols = augmentedDim;
	model->hSzMT.numRows = numSigmaPoints;

	/*
	 Sz = scaleFactor * [predictionCov               0
	 0             processNoiseCov]

	 hSzM = [ state  (state + Sz)  (state - Sz) ]'
	 */
	/*set first initial sigma point*/
	for (col = 0; col < augmentedDim; col++) {
		if (col < stateDim) {
			stateElement = state[col];
			sigmaPointsArray[col] = stateElement;
		} else {
			sigmaPointsArray[col] = 0;
		}
	}
	for (row = 0; row < processCovDim; row++) {
		stateElement = state[row];
		t1 = row * processCovDim;
		for (col = 1; col < augmentedDimAdd; col++) {
			arrIndex = col * augmentedDim + row;
			if (col < processCovDimAdd) {
				covarianceElement = (scaleFactor * predictedCov[t1 + (col - 1)]);
				sigmaPointsArray[arrIndex] = stateElement + covarianceElement;
				sigmaPointsArrayAug[arrIndex] = stateElement - covarianceElement;
			} else {
				sigmaPointsArray[arrIndex] = stateElement;
				sigmaPointsArrayAug[arrIndex] = stateElement;
			}

		}
	}
	/*Fill the covariance*/
	for (row = processCovDim; row < augmentedDim; row++) {
		t2 = (row - processCovDim) * processNoiseCovDim;
		for (col = processCovDimAdd; col < augmentedDimAdd; col++) {
			arrIndex = col * augmentedDim + row;
			covarianceElement = scaleFactor * processCov[t2 + (col - processCovDim - 1)];
			sigmaPointsArray[arrIndex] = covarianceElement;
			sigmaPointsArrayAug[arrIndex] = -covarianceElement;
		}
	}
	/*Fill the rest with 0's*/
	for (col = 1; col < processCovDimAdd; col++) {
		t3 = col * augmentedDim;
		for (row = processCovDim; row < augmentedDim; row++) {
			arrIndex = t3 + row;
			sigmaPointsArray[arrIndex] = 0.0f;
			sigmaPointsArrayAug[arrIndex] = 0.0f;
		}
	}
	return 1;
}
int ukf_calculatePropagatedCovariance(ukf_predictionModel *model, arm_matrix_instance_f32 * propSigmaPoints, arm_matrix_instance_f32 * calculatedCovariance, uint8_t enableCalcMeasurementPt) {
	const float32_t weightCov0 = model->weightCov0;
	const float32_t weightCov1 = model->weightCov1;

	const int xDim = model->stateDim;
	const int numRows = propSigmaPoints->numRows;
	const int numCols = propSigmaPoints->numCols;
	const int numColsOutput = numRows - 1;
	const int L = (numRows - 1) >> 1;
	const int sywTotalCols = model->Syw.numCols;

	float32_t *sigmaPointsData = propSigmaPoints->pData;
	float32_t *outputData = calculatedCovariance->pData;
	float32_t *syxData = model->Syx.pData;
	float32_t *sywData = model->Syw.pData;
	float32_t *syxTData = model->SyxT.pData;
	float32_t *dData = model->D.pData;

	model->D.numRows = numCols;
	model->D.numCols = L;
	calculatedCovariance->numRows = numCols;
	calculatedCovariance->numCols = numColsOutput;

	int row, col, tmp1, tmp2, tmp3;
	float32_t data_row_plus_1_col;
	float32_t data_row_plus_1_plus_L_col;
	float32_t cElement;
	float32_t dElement;

	if (enableCalcMeasurementPt) {
		for (row = 0; row < L; row++) {
			tmp1 = (row + 1) * numCols;
			tmp2 = tmp1 + L*numCols;
			for (col = 0; col < numCols; col++) {
				// cache data. data_row_col
				data_row_plus_1_col = sigmaPointsData[tmp1+ col];
				data_row_plus_1_plus_L_col = sigmaPointsData[tmp2 + col];
				cElement = weightCov0 * (data_row_plus_1_col - data_row_plus_1_plus_L_col);
				dElement = weightCov1 * (data_row_plus_1_col + data_row_plus_1_plus_L_col - (2.0f * sigmaPointsData[col]));
				//dElement = 1000000.0f;
				// transpose and assign at the same time. reversing row and col
				tmp3 = col * numColsOutput + row;
				// output[col][row]
				outputData[tmp3] = cElement;
				// output[col][row+L]
				outputData[tmp3 + L] = dElement;
				if (row < xDim) {
					syxData[col * xDim + row] = cElement;
					syxTData[row * sywTotalCols + col] = cElement;
				} else if (row < xDim + numCols) {
					sywData[col * sywTotalCols + row - xDim] = cElement;
				}
				dData[col * L + row] = dElement;
			}
		}
	} else {
		for (row = 0; row < L; row++) {
			tmp1 = (row + 1) * numCols;
			tmp2 = tmp1 + L*numCols;
			for (col = 0; col < numCols; col++) {
				// cache data. data_row_col
				data_row_plus_1_col = sigmaPointsData[tmp1 + col];
				data_row_plus_1_plus_L_col = sigmaPointsData[tmp2 + col];

				cElement = weightCov0 * (data_row_plus_1_col - data_row_plus_1_plus_L_col);
				dElement = weightCov1 * ((data_row_plus_1_col + data_row_plus_1_plus_L_col) - (2.0f * sigmaPointsData[col]));

				// transpose and assign at the same time. reversing row and col
				tmp3 = col * numColsOutput + row;
				// output[col][row]
				outputData[tmp3] = cElement;
				// output[col][row+L]
				outputData[tmp3 + L] = dElement;
			}
		}
	}

	return 1;
}

int ukf_processPrediction(ukf_predictionModel *model, float32_t * exogenousInput, float32_t dt) {
	int sigmaPointsNum;
	float32_t *inState, *outState, *noise, *X, *hSzMT;

	ukf_predictionUpdate * predictionUpdateFx = model->fx_predictionUpdate;

	float32_t * state = model->state.pData;
	const int stateDim = model->stateDim;
	const int noiseDim = model->observationCovDim;
	const int augmentedStateDim = stateDim + noiseDim;

	const float32_t weightSum0 = model->weightSum0;
	const float32_t weightSum1 = model->weightSum1;

	X = (model->X).pData;
	hSzMT = (model->hSzMT).pData;

	getTransposedSigmaPoints(model, model->observationCov);
	debug_mprint("Sigma points", &(model->hSzMT));

	/*Sigma points number, they are rows since our SP is transposed*/
	sigmaPointsNum = model->hSzMT.numRows;

	/*Initial propagation and weights*/
	outState = (model->X).pData;
	inState = (model->hSzMT).pData;
	noise = inState + stateDim;
	predictionUpdateFx(inState, outState, noise, exogenousInput, dt);

	for (int wSum = 0; wSum < stateDim; wSum++) {
		state[wSum] = weightSum0 * outState[wSum];
	}
	/*now do the rest of it..*/
	for (int sp = 1; sp < sigmaPointsNum+1; sp++) {
		outState = X + sp * stateDim;
		inState = hSzMT + sp * augmentedStateDim;
		noise = inState + stateDim;

		predictionUpdateFx(inState, outState, noise, exogenousInput, dt);

		/*Also have to do weights here*/
		for (int wSum = 0; wSum < stateDim; wSum++) {
			state[wSum] += (weightSum1 * outState[wSum]);
		}
		//debug_mprint("State Loop", &(model->state));

	}
	debug_mprint("propagated points", &(model->X));
	debug_mprint("prediction State", &(model->state));

	ukf_calculatePropagatedCovariance(model, &(model->X), &(model->Sxn), FALSE);
	debug_mprint("propagated prediction covariance", &(model->Sxn));

	qrDecompositionT_f32(&(model->Sxn), NULL, &(model->SxnT));
	arm_mat_trans_f32(&model->SxnT, &model->Sx);

	debug_mprint("prediction sx", &(model->Sx));

	return 1;
}

//todo: since we are transposing sx every time we finish, we can use the transposed matrix directly , optimize code for that.
int ukf_processObservation(ukf_predictionModel *predictModel, ukf_observationModel *obsModel, float32_t * exogenousInput) {
	float32_t *inState, *outState, *noise, *Y, *hSzMT;
	float32_t sigmaPointsNum;
	int row, col,t1,t2,t3;

	ukf_observationUpdate * predictionUpdateFx = obsModel->fx_observationUpdate;
	const int stateDim = predictModel->stateDim;
	const int noiseDim = obsModel->noiseDim;
	const int augmentedStateDim = stateDim + noiseDim;
	// 0th weighted sum is different than prediction
	const float32_t weightSum0 = (predictModel->scaleFactor - stateDim - noiseDim) / predictModel->scaleFactor;
	const float32_t weightSum1 = predictModel->weightSum1;

	float32_t * state = (predictModel->state).pData;
	float32_t * s5Data = predictModel->s5.pData;
	float32_t * propagatedObservationState = (predictModel->yh).pData;
	float32_t * innovationVector = (predictModel->innov).pData;
	float32_t * innovationUpdateVector = (predictModel->innovUpdate).pData;
	Y = (predictModel->Y).pData;
	hSzMT = (predictModel->hSzMT).pData;

	/*get sigma points, passing in the observation covariance*/
	getTransposedSigmaPoints(predictModel, obsModel->cov);
	debug_mprint("State", &(predictModel->state));

	debug_mprint("Sigma Points", &(predictModel->hSzMT));

	/*Sigma points number, they are rows since our SP is transposed*/
	sigmaPointsNum = predictModel->hSzMT.numRows;

	/*Set all dimensions*/
	/*propagated result matrix*/
	(predictModel->Y).numRows = sigmaPointsNum;
	(predictModel->Y).numCols = noiseDim;

	/*weighted obs*/
	(predictModel->yh).numRows = noiseDim;

	/*our covariance matrix*/
	(predictModel->Syn).numCols = sigmaPointsNum - 1;
	(predictModel->Syn).numRows = noiseDim;

	(predictModel->SynT).numCols = noiseDim;
	(predictModel->SynT).numRows = noiseDim;

	(predictModel->Sy).numCols = noiseDim;
	(predictModel->Sy).numRows = noiseDim;

	(predictModel->Syx).numRows = noiseDim;
	(predictModel->Syx).numCols = stateDim;

	(predictModel->SyxT).numRows = stateDim;
	(predictModel->SyxT).numCols = noiseDim;

	(predictModel->Syw).numRows = noiseDim;
	(predictModel->Syw).numCols = noiseDim;

	(predictModel->Pxy).numRows = stateDim;
	(predictModel->Pxy).numCols = noiseDim;

	(predictModel->KG).numRows = stateDim;
	(predictModel->KG).numCols = noiseDim;

	(predictModel->innov).numRows = noiseDim;
	(predictModel->innovUpdate).numRows = noiseDim;

	/*end set all state dim*/

	/*Initial propagation and weights*/
	outState = (predictModel->Y).pData;
	inState = (predictModel->hSzMT).pData;
	noise = inState + stateDim;
	predictionUpdateFx(inState, outState, noise, exogenousInput);

	for (int wSum = 0; wSum < noiseDim; wSum++) {
		propagatedObservationState[wSum] = weightSum0 * outState[wSum];
	}
	/*now do the rest of it..*/
	for (int sp = 1; sp < sigmaPointsNum; sp++) {
		outState = Y + sp * noiseDim;
		inState = hSzMT + sp * augmentedStateDim;
		noise = inState + stateDim;

		predictionUpdateFx(inState, outState, noise, exogenousInput);

		/*Also have to do weights here*/
		for (int wSum = 0; wSum < noiseDim; wSum++) {
			propagatedObservationState[wSum] += weightSum1 * outState[wSum];
		}
	}
	debug_mprint("propagated points", &(predictModel->Y));

	ukf_calculatePropagatedCovariance(predictModel, &(predictModel->Y), &(predictModel->Syn), TRUE);

	debug_mprint("cov", &(predictModel->Syn));
	debug_mprint("syx", &(predictModel->Syx));
	debug_mprint("syw", &(predictModel->Syw));

	qrDecompositionT_f32(&(predictModel->Syn), NULL, &(predictModel->SynT));
	arm_mat_trans_f32(&predictModel->SynT, &predictModel->Sy);

	/*kalman gain steps*/

	debug_mprint("SX", &(predictModel->Sx));
	debug_mprint("SyxT", &(predictModel->SyxT));

	arm_mat_mult_f32(&(predictModel->Sx), &(predictModel->SyxT), &(predictModel->Pxy));
	debug_mprint("pxy", &(predictModel->Pxy));

	/* KG = (Pxy / SyT) / Sy */
	matrixDiv_f32(&(predictModel->KG), &(predictModel->Pxy), &predictModel->SynT, &predictModel->s1, &predictModel->s2, &predictModel->s3);
	matrixDiv_f32(&(predictModel->KG), &(predictModel->KG), &predictModel->Sy, &predictModel->s1, &predictModel->s2, &predictModel->s3);

	debug_mprint("KG", &(predictModel->KG));
	debug_mprint("yh", &(predictModel->yh));

	for (int iv = 0; iv < noiseDim; iv++) {
		innovationVector[iv] = exogenousInput[iv] - propagatedObservationState[iv];
	}
	debug_mprint("innovationUpdate", &(predictModel->innov));

	arm_mat_mult_f32(&(predictModel->KG), &(predictModel->innov), &(predictModel->innovUpdate));

	for (int inup = 0; inup < stateDim; inup++) {
		state[inup] = state[inup] + innovationUpdateVector[inup];
	}
	debug_mprint("newstate", &(predictModel->state));

	predictModel->s1.numRows = stateDim;
	predictModel->s1.numCols = stateDim;

	predictModel->s3.numRows = stateDim;
	predictModel->s3.numCols = noiseDim;

	predictModel->s4.numRows = stateDim;
	predictModel->s4.numCols = stateDim + noiseDim;

	debug_mprint("D", &(predictModel->D));

	/* s1 = KG * Syx */
	arm_mat_mult_f32(&(predictModel->KG), &(predictModel->Syx), &(predictModel->s1));
	debug_mprint("KG * Syx", &(predictModel->s1));

	/* s3 = KG * Syw */
	arm_mat_mult_f32(&(predictModel->KG), &(predictModel->Syw), &(predictModel->s3));
	debug_mprint("KG * Syw", &(predictModel->s3));

	/* s4 = KG * D */
	arm_mat_mult_f32(&(predictModel->KG), &(predictModel->D), &(predictModel->s4));
	debug_mprint("kg*d", &(predictModel->s4));


	/* s5 = [s2 s3 s4] */
	const int s1Rows = predictModel->s1.numRows;
	const int s1Cols = predictModel->s1.numCols;
	const int s3Cols = predictModel->s3.numCols;
	const int s4Cols = predictModel->s4.numCols;
	const int s5Cols = s1Cols + s3Cols + s4Cols;

	predictModel->s5.numRows = s1Rows;
	predictModel->s5.numCols = s5Cols;

	/* copy s2 */
	for (row = 0; row < s1Rows; row++) {
		t1 =row * s5Cols;
		t2 = row * s1Cols;
		for (col = 0; col < s1Cols; col++) {
			t3 = t2 + col;
			s5Data[t1 + col] = predictModel->Sx.pData[t3] - predictModel->s1.pData[t3];
		}
	}

	/* copy s3 */
	for (row = 0; row < s1Rows; row++) {
		t1 = row * s5Cols + s1Cols;
		t2 = row * s3Cols;
		for(col = 0; col < s3Cols; col++){
			s5Data[t1 + col] = predictModel->s3.pData[t2 + col];
		}
	}

	/* copy s4 */
	for (row = 0; row < s1Rows; row++) {
		t1 = row * s5Cols + s1Cols + s3Cols;
		t2 = row * s4Cols;
		for (col = 0; col < s4Cols; col++) {
			s5Data[t1 + col] = predictModel->s4.pData[t2 + col];
		}
	}

	debug_mprint("Pre QR", &(predictModel->s5));

	qrDecompositionT_f32(&(predictModel->s5), NULL, &(predictModel->SxnT));
	arm_mat_trans_f32(&predictModel->SxnT, &predictModel->Sx);

	debug_mprint("final sx", &(predictModel->Sx));

	return 1;
}
